package spark

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

/**
  * @author pinker on 2018/3/15
  */
class MyVariance(mean: Double) extends UserDefinedAggregateFunction {

  override def inputSchema: StructType = StructType(StructField("age", DoubleType, true) :: Nil)

  override def bufferSchema: StructType = StructType(StructField("variance", DoubleType) :: StructField("count", LongType) :: Nil)

  override def dataType: DataType = DoubleType

  override def deterministic: Boolean = true

  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0.0
    buffer(1) = 0L
  }

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)) {
      println("age is :" + input.getDouble(0))
      buffer(0) = buffer.getDouble(0) + (input.getDouble(0) - mean) * (input.getDouble(0) - mean)
      buffer(1) = buffer.getLong(1) + 1
    }
  }

  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  override def evaluate(buffer: Row): Double = {
    if (buffer.getLong(1) > 1) {
      return buffer.getDouble(0) / (buffer.getLong(1) - 1)
    }
    buffer.getDouble(0)
  }
}


